#ifndef BPLUSTREE_H
#define BPLUSTREE_H

#include <iostream>
#include <vector>

template<typename T>
struct BPlusNode {
    std::vector<T> keys;
    std::vector<BPlusNode*> children;
    BPlusNode* next; // 叶子节点链表1
    bool isLeaf;
    int numKeys;
    
    BPlusNode(bool leaf = true) : isLeaf(leaf), numKeys(0), next(nullptr) {}
};

template<typename T>
class BPlusTree {
private:
    BPlusNode<T>* root;
    BPlusNode<T>* leafHead; // 叶子节点链表的头
    int t; // 最小度数
    
    void insertNonFull(BPlusNode<T>* node, T key);
    void splitChild(BPlusNode<T>* parent, int index, BPlusNode<T>* child);
    bool searchHelper(BPlusNode<T>* node, T key) const;
    void removeHelper(BPlusNode<T>* node, T key);
    void removeFromLeaf(BPlusNode<T>* node, int index);
    void removeFromNonLeaf(BPlusNode<T>* node, int index);
    T getPred(BPlusNode<T>* node, int index);
    T getSucc(BPlusNode<T>* node, int index);
    void fill(BPlusNode<T>* node, int index);
    void borrowFromPrev(BPlusNode<T>* node, int index);
    void borrowFromNext(BPlusNode<T>* node, int index);
    void merge(BPlusNode<T>* node, int index);
    void printHelper(BPlusNode<T>* node, int level) const;
    void clearHelper(BPlusNode<T>* node);
    void printLeaves() const;
    
public:
    BPlusTree(int degree = 3);
    ~BPlusTree();
    void insert(T key);
    void remove(T key);
    bool search(T key) const;
    void print() const;
    void printLeafList() const;
    void clear();
};

template<typename T>
BPlusTree<T>::BPlusTree(int degree) : t(degree) {
    root = new BPlusNode<T>(true);
    leafHead = root;
}

template<typename T>
BPlusTree<T>::~BPlusTree() {
    clear();
}

template<typename T>
void BPlusTree<T>::insert(T key) {
    BPlusNode<T>* r = root;
    
    if (r->numKeys == 2 * t - 1) {
        BPlusNode<T>* s = new BPlusNode<T>(false);
        root = s;
        s->children.push_back(r);
        splitChild(s, 0, r);
        insertNonFull(s, key);
    } else {
        insertNonFull(r, key);
    }
}

template<typename T>
void BPlusTree<T>::insertNonFull(BPlusNode<T>* node, T key) {
    int i = node->numKeys - 1;
    
    if (node->isLeaf) {
        while (i >= 0 && node->keys[i] > key) {
            node->keys.insert(node->keys.begin() + i + 1, node->keys[i]);
            i--;
        }
        node->keys.insert(node->keys.begin() + i + 1, key);
        node->numKeys++;
    } else {
        while (i >= 0 && node->keys[i] > key) {
            i--;
        }
        i++;
        
        if (node->children[i]->numKeys == 2 * t - 1) {
            splitChild(node, i, node->children[i]);
            if (node->keys[i] < key) {
                i++;
            }
        }
        insertNonFull(node->children[i], key);
    }
}

template<typename T>
void BPlusTree<T>::splitChild(BPlusNode<T>* parent, int index, BPlusNode<T>* child) {
    BPlusNode<T>* newNode = new BPlusNode<T>(child->isLeaf);
    
    // 移动后半部分键值到新节点
    for (int j = 0; j < t - 1; j++) {
        newNode->keys.push_back(child->keys[j + t]);
    }
    
    // 如果不是叶子节点，移动后半部分子节点
    if (!child->isLeaf) {
        for (int j = 0; j < t; j++) {
            newNode->children.push_back(child->children[j + t]);
        }
    }
    
    // 调整原节点的键值数量
    child->numKeys = t - 1;
    newNode->numKeys = t - 1;
    
    // 在父节点中插入新键值
    if (child->isLeaf) {
        // B+树中，叶子节点分裂时，中间键值会复制到父节点
        parent->keys.insert(parent->keys.begin() + index, child->keys[t - 1]);
        // 更新叶子节点链表
        newNode->next = child->next;
        child->next = newNode;
    } else {
        // 非叶子节点分裂时，中间键值会提升到父节点
        parent->keys.insert(parent->keys.begin() + index, child->keys[t - 1]);
    }
    
    parent->children.insert(parent->children.begin() + index + 1, newNode);
    parent->numKeys++;
    
    // 从原节点中移除已移动的键值
    child->keys.resize(t - 1);
    if (!child->isLeaf) {
        child->children.resize(t);
    }
}

template<typename T>
bool BPlusTree<T>::search(T key) const {
    return searchHelper(root, key);
}

template<typename T>
bool BPlusTree<T>::searchHelper(BPlusNode<T>* node, T key) const {
    int i = 0;
    while (i < node->numKeys && key > node->keys[i]) {
        i++;
    }
    
    if (node->isLeaf) {
        return (i < node->numKeys && key == node->keys[i]);
    }
    
    return searchHelper(node->children[i], key);
}

template<typename T>
void BPlusTree<T>::remove(T key) {
    if (!root) return;
    
    removeHelper(root, key);
    
    if (root->numKeys == 0) {
        BPlusNode<T>* oldRoot = root;
        if (root->isLeaf) {
            root = nullptr;
            leafHead = nullptr;
        } else {
            root = root->children[0];
        }
        delete oldRoot;
    }
}

template<typename T>
void BPlusTree<T>::removeHelper(BPlusNode<T>* node, T key) {
    int index = 0;
    while (index < node->numKeys && node->keys[index] < key) {
        index++;
    }
    
    if (node->isLeaf) {
        if (index < node->numKeys && node->keys[index] == key) {
            removeFromLeaf(node, index);
        }
    } else {
        if (index < node->numKeys && node->keys[index] == key) {
            removeFromNonLeaf(node, index);
        } else {
            if (node->children[index]->numKeys < t) {
                fill(node, index);
            }
            
            if (index > node->numKeys) {
                removeHelper(node->children[index - 1], key);
            } else {
                removeHelper(node->children[index], key);
            }
        }
    }
}

template<typename T>
void BPlusTree<T>::removeFromLeaf(BPlusNode<T>* node, int index) {
    for (int i = index + 1; i < node->numKeys; i++) {
        node->keys[i - 1] = node->keys[i];
    }
    node->numKeys--;
}

template<typename T>
void BPlusTree<T>::removeFromNonLeaf(BPlusNode<T>* node, int index) {
    T key = node->keys[index];
    
    if (node->children[index]->numKeys >= t) {
        T pred = getPred(node, index);
        node->keys[index] = pred;
        removeHelper(node->children[index], pred);
    } else if (node->children[index + 1]->numKeys >= t) {
        T succ = getSucc(node, index);
        node->keys[index] = succ;
        removeHelper(node->children[index + 1], succ);
    } else {
        merge(node, index);
        removeHelper(node->children[index], key);
    }
}

template<typename T>
T BPlusTree<T>::getPred(BPlusNode<T>* node, int index) {
    BPlusNode<T>* current = node->children[index];
    while (!current->isLeaf) {
        current = current->children[current->numKeys];
    }
    return current->keys[current->numKeys - 1];
}

template<typename T>
T BPlusTree<T>::getSucc(BPlusNode<T>* node, int index) {
    BPlusNode<T>* current = node->children[index + 1];
    while (!current->isLeaf) {
        current = current->children[0];
    }
    return current->keys[0];
}

template<typename T>
void BPlusTree<T>::fill(BPlusNode<T>* node, int index) {
    if (index != 0 && node->children[index - 1]->numKeys >= t) {
        borrowFromPrev(node, index);
    } else if (index != node->numKeys && node->children[index + 1]->numKeys >= t) {
        borrowFromNext(node, index);
    } else {
        if (index != node->numKeys) {
            merge(node, index);
        } else {
            merge(node, index - 1);
        }
    }
}

template<typename T>
void BPlusTree<T>::borrowFromPrev(BPlusNode<T>* node, int index) {
    BPlusNode<T>* child = node->children[index];
    BPlusNode<T>* sibling = node->children[index - 1];
    
    for (int i = child->numKeys - 1; i >= 0; i--) {
        child->keys[i + 1] = child->keys[i];
    }
    
    if (!child->isLeaf) {
        for (int i = child->numKeys; i >= 0; i--) {
            child->children[i + 1] = child->children[i];
        }
    }
    
    child->keys[0] = node->keys[index - 1];
    
    if (!child->isLeaf) {
        child->children[0] = sibling->children[sibling->numKeys];
    }
    
    node->keys[index - 1] = sibling->keys[sibling->numKeys - 1];
    
    child->numKeys++;
    sibling->numKeys--;
}

template<typename T>
void BPlusTree<T>::borrowFromNext(BPlusNode<T>* node, int index) {
    BPlusNode<T>* child = node->children[index];
    BPlusNode<T>* sibling = node->children[index + 1];
    
    child->keys[child->numKeys] = node->keys[index];
    
    if (!child->isLeaf) {
        child->children[child->numKeys + 1] = sibling->children[0];
    }
    
    node->keys[index] = sibling->keys[0];
    
    for (int i = 1; i < sibling->numKeys; i++) {
        sibling->keys[i - 1] = sibling->keys[i];
    }
    
    if (!sibling->isLeaf) {
        for (int i = 1; i <= sibling->numKeys; i++) {
            sibling->children[i - 1] = sibling->children[i];
        }
    }
    
    child->numKeys++;
    sibling->numKeys--;
}

template<typename T>
void BPlusTree<T>::merge(BPlusNode<T>* node, int index) {
    BPlusNode<T>* child = node->children[index];
    BPlusNode<T>* sibling = node->children[index + 1];
    
    child->keys[child->numKeys] = node->keys[index];
    
    for (int i = 0; i < sibling->numKeys; i++) {
        child->keys[child->numKeys + 1 + i] = sibling->keys[i];
    }
    
    if (!child->isLeaf) {
        for (int i = 0; i <= sibling->numKeys; i++) {
            child->children[child->numKeys + 1 + i] = sibling->children[i];
        }
    } else {
        // 更新叶子节点链表
        child->next = sibling->next;
    }
    
    for (int i = index + 1; i < node->numKeys; i++) {
        node->keys[i - 1] = node->keys[i];
    }
    
    for (int i = index + 2; i <= node->numKeys; i++) {
        node->children[i - 1] = node->children[i];
    }
    
    child->numKeys += sibling->numKeys + 1;
    node->numKeys--;
    
    delete sibling;
}

template<typename T>
void BPlusTree<T>::print() const {
    if (root) {
        printHelper(root, 0);
    }
    std::cout << std::endl;
}

template<typename T>
void BPlusTree<T>::printHelper(BPlusNode<T>* node, int level) const {
    std::cout << "Level " << level << ": ";
    for (int i = 0; i < node->numKeys; i++) {
        std::cout << node->keys[i] << " ";
    }
    std::cout << std::endl;
    
    if (!node->isLeaf) {
        for (int i = 0; i <= node->numKeys; i++) {
            printHelper(node->children[i], level + 1);
        }
    }
}

template<typename T>
void BPlusTree<T>::printLeafList() const {
    std::cout << "Leaf List: ";
    BPlusNode<T>* current = leafHead;
    while (current) {
        for (int i = 0; i < current->numKeys; i++) {
            std::cout << current->keys[i] << " ";
        }
        current = current->next;
    }
    std::cout << std::endl;
}

template<typename T>
void BPlusTree<T>::clear() {
    clearHelper(root);
    root = nullptr;
    leafHead = nullptr;
}

template<typename T>
void BPlusTree<T>::clearHelper(BPlusNode<T>* node) {
    if (node) {
        if (!node->isLeaf) {
            for (int i = 0; i <= node->numKeys; i++) {
                clearHelper(node->children[i]);
            }
        }
        delete node;
    }
}

#endif 